package com.sa.gateway.config;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

/**
 * @author Architecture Station
 * @version 1.0
 */
@Configuration
public class RateLimiterConfig {

    @Autowired
    private ReactiveStringRedisTemplate redisTemplate;

    @Autowired
    private RedisScript<List<Long>> script;

    @Autowired
    private ConfigurationService configurationService;

    private static final Map<String, Integer> REPLENISH_RATES = new HashMap<>();
    private static final Map<String, Integer> BURST_CAPACITIES = new HashMap<>();

    static {
        // 根据用户类型设置不同的限流参数
        REPLENISH_RATES.put("Normal", 10); // 针对普通用户，令牌每秒产生的速度是10个令牌
        BURST_CAPACITIES.put("Normal", 20); // 针对普通用户，令牌桶容量是20个令牌
        REPLENISH_RATES.put("VIP", 50); // 针对VIP用户，令牌每秒产生的速度是50个令牌
        BURST_CAPACITIES.put("VIP", 100); // 针对VIP用户，令牌桶容量是100个令牌
    }

    @Bean
    @Primary
    public RateLimiter<RedisRateLimiter.Config> dynamicRateLimiter() {
        return new RedisRateLimiter(redisTemplate, script, configurationService) {
            @Override
            public Mono<Response> isAllowed(String routeId, String id) {
                String[] parts = id.split(":");
                String userType = parts[0];
                int replenishRate = REPLENISH_RATES.getOrDefault(userType, 10);
                int burstCapacity = BURST_CAPACITIES.getOrDefault(userType, 20);

                // Create a new configuration for the rate limiter
                Config config = new Config().setReplenishRate(replenishRate).setBurstCapacity(burstCapacity);

                // Get the configuration from the cache or create a new one
                Config existingConfig = getConfig().get(routeId);
                if (existingConfig == null || existingConfig.getReplenishRate() != replenishRate || existingConfig.getBurstCapacity() != burstCapacity) {
                    getConfig().put(routeId, config);
                }

                return super.isAllowed(routeId, id);
            }
        };
    }

    @Bean
    public KeyResolver userKeyResolver() {
        return exchange -> Mono.just(
                Optional.ofNullable(exchange.getRequest().getHeaders().getFirst("X-User-Type"))
                        .orElse("Normal")
        );
    }
}
